# 逆运动学这里主要方法总结为：
# Pseudo inversing a Jacobian iteratively until convergence
# v_star求法有变
import pinocchio as pin
from pinocchio.utils import rand, zero
import meshcat
import numpy as np
# 引入src中所有的函数
import sys
from os.path import dirname, join

sys.path.append(join(dirname(dirname(__file__)), "src"))
sys.path.append(dirname(dirname(__file__))) # 把文件根路径添加到sys.path中，这样在根路径下的src文件夹才能被import
from src import *
# <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<


# 获取URDF模型路径
URDF_FILE_PATH = filePath(URDF_NAME = 'ur5_gripper', ROBOT = True)
# <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

# 构建robot对象，并得到viz
robot, viz = robot_viz(URDF_FILE_PATH)
q0 = pin.randomConfiguration(robot.model)
pin.framesForwardKinematics(robot.model, robot.data, q0)
pin.updateFramePlacements(robot.model, robot.data) 
robot.initViewer(loadModel=True)
viz.display(q0)
# <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
# <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
import time
import math
# 需要求伪逆
from numpy.linalg import pinv
# 设置初始位置
q = q0
idx = 6
print('idx is ',idx)
dt_ik = 0.001
ik_time_max = 5.0
IT_MAX = math.floor(ik_time_max / dt_ik)
damp = 1e-12
##########################################



# <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
# 开始进行逆运动学求解


# 设置effector目标位置和姿态
T_t = pin.SE3(np.eye(3), np.array([0.5, 0.3, 0.0]))
homo_matrix = T_t.homogeneous
############################################################
# 创建一个box位于目标位置，方便查看
box_name = "world/box"
box_size = [0.04, 0.02, 0.01]

viz.viewer[box_name].set_object(
    meshcat.geometry.Box(box_size),
    meshcat.geometry.MeshLambertMaterial(color=0x6622cc,reflectivity=0.4)
)

viz.viewer[box_name].set_transform(
    homo_matrix
)
# <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<


i = 0
while True:
    
    
    T_bt = robot.data.oMi[idx].actInv(T_t)
    J_b = pin.computeJointJacobian(robot.model, robot.data, q, idx)
    J_l = -pin.Jlog6(T_bt.inverse())

    # 接下来会这里会利用J_error v^{*} = -e(q)来求v^{*}
    J_error = J_l.dot(J_b) # 误差Jacobian
    e_q = pin.log(T_bt).vector # 定义的误差向量
    print('e_q is ',e_q)
    print('J_error is ',J_error)
    # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
    v_star = pinv(J_error).dot(-e_q) # 此处是求解的关键，这里直接用pinv，没有借用damping项
    # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
    q = pin.integrate(robot.model, q, v_star*dt_ik)
    pin.forwardKinematics(robot.model, robot.data, q)# 更新位置，这个一定不能忘
    viz.display(q)

    if np.linalg.norm(e_q) < 1e-4:
        success = True
        break
    if i>IT_MAX:
        success = False
        break
    i += 1


if success:
    print("Convergence achieved!")
else:
    print("Convergence failed!")

print('Final q is ',q)
print(robot.data.oMi[idx].rotation)
from scipy.spatial.transform import Rotation as R
quaternion = R.from_matrix(robot.data.oMi[idx].rotation).as_quat()
print('Final quaternion is ',quaternion)

viz.display(q)